import torch

class LastFrameExtractorNode:
    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
            },
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("last_image",)
    FUNCTION = "extract_last_frame"
    CATEGORY = "image"
    
    def extract_last_frame(self, images):
        """
        Extract the last image from a batch of images.
        
        Args:
            images: Input image tensor with shape [batch_size, height, width, channels]
        
        Returns:
            Tuple containing the last image as a tensor with shape [1, height, width, channels]
        """
        if images is None or len(images) == 0:
            raise ValueError("No images provided")
        
        # Get the last image from the batch
        last_image = images[-1:].clone()  # Keep batch dimension by using [-1:]
        
        return (last_image,)

# Required mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
    "LastFrameExtractorNode": LastFrameExtractorNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LastFrameExtractorNode": "Last Frame Extractor"
}